package cn.org.itool.ws.framework.protocol.pomelo;

import lombok.Data;

import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public class Protocol {

    private static final int PKG_HEAD_BYTES = 4;
    private static final int MSG_FLAG_BYTES = 1;
    private static final int MSG_ROUTE_CODE_BYTES = 2;
    private static final int MSG_ID_MAX_BYTES = 5;
    private static final int MSG_ROUTE_LEN_BYTES = 1;

    private static final int MSG_ROUTE_CODE_MAX = 0xffff;
    private static final int MSG_COMPRESS_ROUTE_MASK = 0x1;
    private static final int MSG_COMPRESS_GZIP_MASK = 0x1;

    private static final int MSG_COMPRESS_GZIP_ENCODE_MASK = 1 << 4;
    private static final int MSG_TYPE_MASK = 0x7;

    public static byte[] strEncode(String str) {
        return str.getBytes(StandardCharsets.UTF_8);
    }

    public static String strDecode(byte[] buffer) {
        int offset = 0;
        int end = buffer.length;
        List<Character> array = new ArrayList<>();
        int charCode;
        while (offset < end) {
            if ((buffer[offset] & 0xFF) < 128) {
                charCode = buffer[offset];
                offset += 1;
            } else if ((buffer[offset] & 0xFF) < 224) {
                charCode = ((buffer[offset] & 0x1f) << 6) + (buffer[offset + 1] & 0x3f);
                offset += 2;
            } else {
                charCode = ((buffer[offset] & 0x0f) << 12) + ((buffer[offset + 1] & 0x3f) << 6) + (buffer[offset + 2] & 0x3f);
                offset += 3;
            }
            array.add((char) charCode);
        }
        StringBuilder sb = new StringBuilder();
        for (Character c : array) {
            sb.append(c);
        }
        return sb.toString();
    }


    /**
     * Pomelo package format:
     * +------+-------------+------------------+
     * | type | body length |       body       |
     * +------+-------------+------------------+
     *
     * Head: 4bytes
     *   0: package type,
     *      1 - handshake,
     *      2 - handshake ack,
     *      3 - heartbeat,
     *      4 - data
     *      5 - kick
     *   1 - 3: big-endian body length
     * Body: body length bytes
     */
    public static class Package {

        /**
         *  客户端到服务器的握手请求以及服务器到客户端的握手响应
         */
        public static final int TYPE_HANDSHAKE = 1;

        /**
         * 客户端到服务器的握手ack
         */
        public static final int TYPE_HANDSHAKE_ACK = 2;

        /**
         * 心跳包
         */
        public static final int TYPE_HEARTBEAT = 3;

        /**
         * 数据包
         */
        public static final int TYPE_DATA = 4;

        /**
         * 服务器主动断开连接通知
         */
        public static final int TYPE_KICK = 5;

        public static byte[] encode(int type, byte[] body) {
            int length = body != null ? body.length : 0;
            byte[] buffer = new byte[PKG_HEAD_BYTES + length];
            int index = 0;
            buffer[index++] = (byte)(type & 0xff);
            buffer[index++] = (byte)((length >> 16) & 0xff);
            buffer[index++] = (byte)((length >> 8) & 0xff);
            buffer[index++] = (byte)(length & 0xff);
            if (body != null) {
                System.arraycopy(body, 0, buffer, index, length);
            }
            return buffer;
        }

        public static List<PackageType> decode(byte[] buffer, int offset) {
            List<PackageType> rs = new ArrayList<>();
            while (offset < buffer.length) {
                byte type = buffer[offset++];
                int length = ((buffer[offset++] & 0xff) << 16) | ((buffer[offset++] & 0xff) << 8) | (buffer[offset++] & 0xff);
                byte[] body = length > 0 ? new byte[length] : null;
                if (body != null) {
                    System.arraycopy(buffer, offset, body, 0, length);
                }
                offset += length;
                rs.add(new PackageType(type, body));
            }
            return rs;
        }

        public static class PackageType {
            byte type;
            byte[] body;

            public PackageType(byte type, byte[] body) {
                this.type = type;
                this.body = body;
            }

            public byte getType() {
                return type;
            }

            public byte[] getBody() {
                return body;
            }
        }
    }


    @Data
    public static class Message{

        /**
         * 消息id
         * request | response
         */
        private Integer id;

        /**
         * 0～7
         * 0-request; 1-notify; 2-response; 3-push
         */
        private Integer type;

        /**
         * gzip 压缩标志位
         * 0-未压缩; 1-压缩
         */
        private int compressGzip;

        /**
         * route压缩标志位
         * 0-未压缩; 1-压缩
         */
        private int compressRoute;

        /**
         * 路由
         */
        private Object route;

        /**
         * 数据
         */
        private byte[] body;


        /*----------------------------------------------------------------------------*/


        /**
         * request 消息
         */
        public static final int TYPE_REQUEST = 0;

        /**
         * notify 消息
         */
        public static final int TYPE_NOTIFY = 1;

        /**
         * response 消息
         */
        public static final int TYPE_RESPONSE = 2;

        /**
         * push 消息
         */
        public static final int TYPE_PUSH = 3;

        public static byte[] encode(Integer id, int type, boolean compressRoute, Object route, byte[] msg, boolean compressGzip) throws Exception {
            // calculate message max length
            int idBytes = msgHasId(type) ? calculateMsgIdBytes(id) : 0;
            int msgLen = MSG_FLAG_BYTES + idBytes;

            if (msgHasRoute(type)) {
                if (compressRoute) {
                    if (!(route instanceof Integer)) {
                        throw new IllegalArgumentException("error flag for number route!");
                    }
                    msgLen += MSG_ROUTE_CODE_BYTES;
                } else {
                    msgLen += MSG_ROUTE_LEN_BYTES;
                    if (route != null) {
                        byte[] routeBytes = Protocol.strEncode(route.toString());
                        if (routeBytes.length > 255) {
                            throw new IllegalArgumentException("route maxlength is overflow");
                        }
                        msgLen += routeBytes.length;
                    }
                }
            }

            if (msg != null) {
                msgLen += msg.length;
            }

            byte[] buffer = new byte[msgLen];
            int offset = 0;

            // add flag
            offset = encodeMsgFlag(type, compressRoute, buffer, offset, compressGzip);

            // add message id
            if (msgHasId(type)) {
                offset = encodeMsgId(id, buffer, offset);
            }

            // add route
            if (msgHasRoute(type)) {
                offset = encodeMsgRoute(compressRoute, route, buffer, offset);
            }

            // add body
            if (msg != null) {
                offset = encodeMsgBody(msg, buffer, offset);
            }

            return buffer;
        }

        public static Message decode(byte[] buffer) {
            ByteArrayInputStream bytes = new ByteArrayInputStream(buffer);
            int bytesLen = buffer.length;
            int offset = 0;
            int id = 0;
            String route = null;

            // parse flag
            int flag = buffer[offset++] & 0xff;
            int compressRoute = flag & MSG_COMPRESS_ROUTE_MASK;
            int type = (flag >> 1) & MSG_TYPE_MASK;
            int compressGzip = (flag >> 4) & MSG_COMPRESS_GZIP_MASK;

            // parse id
            if (msgHasId(type)) {
                int m = 0;
                int i = 0;
                do {
                    m = buffer[offset++] & 0xff;
                    id += (m & 0x7f) << (7 * i);
                    i++;
                } while (m >= 128);
            }

            // parse route
            if (msgHasRoute(type)) {
                if (compressRoute != 0) {
                    int high = buffer[offset++] & 0xff;
                    int low = buffer[offset++] & 0xff;
                    route = Integer.toString((high << 8) | low);
                } else {
                    int routeLen = buffer[offset++] & 0xff;
                    if (routeLen != 0) {
                        byte[] routeBytes = new byte[routeLen];
                        System.arraycopy(buffer, offset, routeBytes, 0, routeLen);
                        route = new String(routeBytes, StandardCharsets.UTF_8);
                        offset += routeLen;
                    } else {
                        route = "";
                    }
                }
            }

            // parse body
            int bodyLen = bytesLen - offset;
            byte[] body = new byte[bodyLen];
            System.arraycopy(buffer, offset, body, 0, bodyLen);


            Message msg = new Message();
            msg.id = id;
            msg.type = type;
            msg.compressRoute = compressRoute;
            msg.route = route;
            msg.body = body;
            msg.compressGzip = compressGzip;
            return msg;
        }

        private static boolean msgHasRoute(int type) {
            return type == TYPE_REQUEST || type == TYPE_NOTIFY ||
                    type == TYPE_PUSH;
        }

        private static boolean msgHasId(int type) {
            return type == TYPE_REQUEST || type == TYPE_RESPONSE;
        }

        private static int calculateMsgIdBytes(int id) {
            int len = 0;
            do {
                len += 1;
                id >>= 7;
            } while (id > 0);
            return len;
        }

        public static int encodeMsgFlag(int type, boolean compressRoute, byte[] buffer, int offset, boolean compressGzip) {
            if (type != Message.TYPE_REQUEST && type != Message.TYPE_NOTIFY &&
                    type != Message.TYPE_RESPONSE && type != Message.TYPE_PUSH) {
                throw new IllegalArgumentException("unkonw message type: " + type);
            }

            buffer[offset] = (byte)((type << 1) | (compressRoute ? 1 : 0));

            if (compressGzip) {
                buffer[offset] |= MSG_COMPRESS_GZIP_ENCODE_MASK;
            }

            return offset + MSG_FLAG_BYTES;
        }

        public static int encodeMsgId(int id, byte[] buffer, int offset) {
            do{
                int tmp = id % 128;
                int next = (int)Math.floor(id/128.0);

                if(next != 0){
                    tmp = tmp + 128;
                }
                buffer[offset++] = (byte) tmp;

                id = next;
            } while(id != 0);

            return offset;
        }

        public static int encodeMsgRoute(boolean compressRoute, Object route, byte[] buffer, int offset) throws Exception {
            if (compressRoute) {
                int routeInt = Integer.parseInt(route.toString());
                if(routeInt > MSG_ROUTE_CODE_MAX){
                    throw new Exception("route number is overflow");
                }

                buffer[offset++] = (byte) (((byte)route >> 8) & 0xff);
                buffer[offset++] = (byte) ((byte)route & 0xff);
            } else {
                if(Objects.nonNull(route)) {
                    byte[] routeBytes = Protocol.strEncode(route.toString());
                    buffer[offset++] = (byte) (routeBytes.length & 0xff);
                    System.arraycopy(routeBytes, 0, buffer, offset, routeBytes.length);
                    offset += routeBytes.length;
                } else {
                    buffer[offset++] = 0;
                }
            }

            return offset;
        }

        public static int encodeMsgBody(byte[] msg, byte[] buffer, int offset) {
            System.arraycopy(msg, 0, buffer, offset, msg.length);
            return offset + msg.length;
        };
    }
}



